package com.project.im.common.util;

import cn.hutool.core.codec.Base64;
import cn.hutool.crypto.SecureUtil;
import lombok.SneakyThrows;
import lombok.experimental.UtilityClass;

import javax.crypto.KeyAgreement;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.DHPublicKeySpec;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

@UtilityClass
public class DHUtils {

    /**
     * 非对称加密密钥算法
     */
    private final String KEY_ALGORITHM = "DH";
    /**
     * 公钥
     */
    private final String PUBLIC_KEY = "DHPublicKey";
    /**
     * 私钥
     */
    private final String PRIVATE_KEY = "DHPrivateKey";
    private final static BigInteger P = new BigInteger("ffffffffffffffffc90fdaa22168c234c4c6628b80dc1cd129024e088a67cc74020bbea63b139b22514a08798e3404ddef9519b3cd3a431b302b0a6df25f14374fe1356d6d51c245e485b576625e7ec6f44c42e9a637ed6b0bff5cb6f406b7edee386bfb5a899fa5ae9f24117c4b1fe649286651ece65381ffffffffffffffff", 16);
    private final static BigInteger G = new BigInteger("2", 16);

    /**
     * 初始化密钥对
     */
    @SneakyThrows
    public Map<String, Object> initKey() {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        DHParameterSpec parameterSpec = new DHParameterSpec(P, G);
        keyPairGenerator.initialize(parameterSpec);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        Map<String, Object> keyMap = new HashMap<>(2);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }

    /**
     * 根据远程公钥初始化密钥对
     */
    @SneakyThrows
    public Map<String, Object> initKey(byte[] remoteKey) {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(remoteKey);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        DHPublicKey remotePublicKey = (DHPublicKey) keyFactory.generatePublic(keySpec);
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        keyPairGenerator.initialize(remotePublicKey.getParams());
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        Map<String, Object> keyMap = new HashMap<>(2);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }

    /**
     * 构建密钥
     */
    @SneakyThrows
    public byte[] getSecretKey(byte[] remoteKey, DHPrivateKey privateKey) {
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        BigInteger y = new BigInteger(1, remoteKey);
        DHPublicKeySpec keySpec = new DHPublicKeySpec(y, privateKey.getParams().getP(), privateKey.getParams().getG());
        PublicKey publicKey = keyFactory.generatePublic(keySpec);
        KeyAgreement keyAgree = KeyAgreement.getInstance(keyFactory.getAlgorithm());
        keyAgree.init(privateKey);
        keyAgree.doPhase(publicKey, true);
        return keyAgree.generateSecret();
    }


    /**
     * 构建密钥
     */
    public String getSecretKey(String publicKey, DHPrivateKey privateKey) {
        byte[] remoteKey = Base64.decode(publicKey);
        byte[] secKey = getSecretKey(remoteKey, privateKey);
        return SecureUtil.md5(Base64.encode(secKey));
    }

    /**
     * 获取私钥
     */
    public DHPrivateKey getPrivateKey(Map<String, Object> keyMap) {
        return (DHPrivateKey) keyMap.get(PRIVATE_KEY);
    }

    /**
     * 获取公钥
     */
    public DHPublicKey getPublicKey(Map<String, Object> keyMap) {
        return (DHPublicKey) keyMap.get(PUBLIC_KEY);
    }
}
